#include <iostream>
#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/features2d/features2d.hpp>


#include <g2o/core/sparse_optimizer.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/robust_kernel.h>
#include <g2o/core/robust_kernel_impl.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/solvers/cholmod/linear_solver_cholmod.h>
#include <g2o/types/slam3d/se3quat.h>
#include <g2o/types/sba/types_six_dof_expmap.h>
#include <g2o/types/slam3d/vertex_pointxyz.h>

#include <eigen3/Eigen/Core>

using namespace std;
using namespace cv;

int findCorrespondingPoints(const cv::Mat &img1, const cv::Mat &img2, vector<cv::Point2f>& points1,
                            vector<cv::Point2f> &points2);

double cx=325.5;
double cy=253.5;
double fx=518.0;
double fy=519.0;

int main(int argc, char **argv) {
    
    if(argc!=3){
        cout<<"Usage: ba_example img1, img2"<<endl;
        exit(1);
    }
    
    cv::Mat img1=imread(argv[1]);
    cv::Mat img2=imread(argv[2]);
    
    assert(img1.data&&img2.data);
    
    vector<cv::Point2f> pts1, pts2;
    if(findCorrespondingPoints(img1,img2, pts1, pts2)==false)
    {
        cout<<"匹配點不夠"<<endl;
        return 0;
    }
    cout<<"匹配点足够"<<endl;
    
    g2o::SparseOptimizer optimizer;
    typedef g2o::BlockSolver< g2o::BlockSolverTraits<6,3> > Block;
    Block::LinearSolverType* linearSolver=new g2o::LinearSolverCholmod<Block::PoseMatrixType > (); //选择线性方程求解器
    
    Block* solver_ptr = new Block( unique_ptr<Block::LinearSolverType>(linearSolver) );
    g2o::OptimizationAlgorithmLevenberg* solver = new g2o::OptimizationAlgorithmLevenberg(unique_ptr<Block>(solver_ptr) );
    optimizer.setAlgorithm(solver); //设置迭代算法
    optimizer.setVerbose(false);  //打开调解输出
    
    //添加节点
    //添加位姿节点（个数为关键帧数目）
    for(int i=0; i<2; i++)
    {
        g2o::VertexSE3Expmap *v=new g2o::VertexSE3Expmap(); //位姿节点类型
        v->setId(i);
        if(i==0)
            v->setFixed(true);
        v->setEstimate(g2o::SE3Quat()); //预设位姿信息
        optimizer.addVertex(v);
    }
    
    //添加特征点节点
    //以第一帧为准
    for(size_t i=0; i<pts1.size(); i++)
    {
        g2o::VertexPointXYZ* v=new g2o::VertexPointXYZ();
        v->setId(2+i);  //前两个接点为位姿节点；
        double z=1;
        double x=(pts1[i].x-cx)*z/fx;
        double y=(pts1[i].y-cy)*z/fy;
        v->setMarginalized(true);
        v->setEstimate(Eigen::Vector3d(x,y,z));
        optimizer.addVertex(v);  //添加关键点节点
    }
    
    g2o::CameraParameters *camera=new g2o::CameraParameters(fx, Eigen::Vector2d(cx, cy), 0);
    camera->setId(0);
    optimizer.addParameter(camera);
    
    vector<g2o::EdgeProjectXYZ2UV*> edges;
    for(size_t i=0; i<pts1.size(); i++)
    {
        g2o::EdgeProjectXYZ2UV*  edge = new g2o::EdgeProjectXYZ2UV();
        edge->setVertex( 0, dynamic_cast<g2o::VertexPointXYZ*>   (optimizer.vertex(i+2)) );
        edge->setVertex( 1, dynamic_cast<g2o::VertexSE3Expmap*>     (optimizer.vertex(0)) );
        edge->setMeasurement( Eigen::Vector2d(pts1[i].x, pts1[i].y ) );
        edge->setInformation( Eigen::Matrix2d::Identity() );
        edge->setParameterId(0, 0);
        
        //核函数:保证每条边的误差不会过大
        edge->setRobustKernel(new g2o::RobustKernelHuber() );  
        optimizer.addEdge(edge);
        edges.push_back(edge);
    }
    
    for(size_t i=0; i<pts2.size(); i++)
    {
        g2o::EdgeProjectXYZ2UV* edge=new g2o::EdgeProjectXYZ2UV();
        edge->setVertex(0, dynamic_cast<g2o::VertexPointXYZ*>  (optimizer.vertex(i+2)) );
        edge->setVertex(1, dynamic_cast<g2o::VertexSE3Expmap* > (optimizer.vertex(1)) );
        edge->setMeasurement(Eigen::Vector2d(pts1[i].x , pts2[i].y) );
        edge->setInformation(Eigen::Matrix2d::Identity());
        edge->setParameterId(0,0);
        
        edge->setRobustKernel(new g2o::RobustKernelHuber() );
        optimizer.addEdge(edge);
        edges.push_back(edge);
    }
    
    cout<<"开始优化"<<endl;
    optimizer.setVerbose(true);
    optimizer.initializeOptimization();
    optimizer.optimize(10);
    cout<<"优化结束"<<endl;
    
    g2o::VertexSE3Expmap *v=dynamic_cast<g2o::VertexSE3Expmap*>(optimizer.vertex(1));
    Eigen::Isometry3d pose=v->estimate();
    cout<<"Pose= "<<endl<<pose.matrix()<<endl;
    
    for(size_t i=0; i<pts1.size(); i++)
    {
        g2o::VertexSBAPointXYZ *v=dynamic_cast<g2o::VertexSBAPointXYZ*> (optimizer.vertex(i+2));
        cout<<"vertex id"<<i+2<<" , pos= ";
        Eigen::Vector3d pos=v->estimate();
        cout<<pos(0)<<" , "<<pos(1)<<" , "<<pos(2)<<endl;
    }
    
    int inliers =0;
    for(auto e:edges)
    {
        e->computeError();
        //chi2 为误差， 如果误差太大，说明这条边与其他边不相符；
        if(e->chi2()>1)
        {
            cout<<"error = "<<e->chi2()<<endl;
        }
        else
        {
            inliers++;
        }
    }
    
    cout<<"inliers in total points: "<<inliers<<"/"<<pts1.size()+pts2.size()<<endl;
    optimizer.save("ba.g2o");
    return 0;
}

int findCorrespondingPoints(const cv::Mat& img1, const cv::Mat& img2, vector<cv::Point2f>& points1, vector<cv::Point2f>& points2)
{
    //cv::ORB orb1;
    Ptr<ORB> detector = ORB::create(100);
    vector<cv::KeyPoint> kp1,kp2;
    cv::Mat desp1, desp2;
    detector->detectAndCompute(img1, cv::Mat(), kp1, desp1);
    detector->detectAndCompute(img2, cv::Mat(), kp2, desp2);
    cout<<"分别找到"<<kp1.size()<<"和"<<kp2.size()<<"个特征点"<<endl;
    
    cv::Ptr<cv::DescriptorMatcher> matcher=cv::DescriptorMatcher::create("BruteForce-Hamming");
    
    double knn_match_ratio=0.8;
    vector<vector<cv::DMatch>> matches_knn;
    matcher->knnMatch(desp1, desp2, matches_knn, 2); //每个点会产生2个匹配点
    vector<cv::DMatch> matches;
    for(size_t i=0; i<matches_knn.size(); i++)
    {
        if(matches_knn[i][0].distance<knn_match_ratio*matches_knn[i][1].distance)
        {
            matches.push_back(matches_knn[i][0]);
        }
    }
    
    if(matches.size()<=20)
        return false;
    
    for(auto m:matches)
    {
        points1.push_back(kp1[m.queryIdx].pt);
        points2.push_back(kp2[m.trainIdx].pt);
    }
    
    return true;
}
    
